#!/bin/bash

# ==============================================
# 目前脚本仅支持ubuntu系统
# AMD vllm openai 通过API_KEY安装及测试
# ==============================================


# ===== 安装ROCm（如果未安装）=====
if ! command -v rocminfo &> /dev/null; then
    echo "正在安装ROCm..."
    wget https://gitee.com/huangsheng32/AMDGPU_linux/raw/master/rocm_install.sh
    sudo chmod +x rocm_install.sh && ./rocm_install.sh
    sudo rm -f rocm_install.sh
    echo "ROCm安装完成！"
else
    echo "ROCm已安装，跳过安装步骤。"
    echo "当前ROCm版本：$(rocminfo | grep 'ROCk version' || echo '未知')"
fi

# ===== 安装Docker（如果未安装）=====
if ! command -v docker &> /dev/null; then
    echo "正在安装Docker..."
    wget https://gitee.com/huangsheng32/linux/raw/master/install_docker.sh
    sudo chmod +x install_docker.sh && ./install_docker.sh
    sudo rm -f install_docker.sh
    echo "Docker安装完成！"
else
    echo "Docker已安装，跳过安装步骤。"
    echo "当前Docker版本：$(docker --version)"
fi

if docker images | grep -q "rcatrangiu/vllm-openai-rocm"; then
    echo "Docker镜像 rcatrangiu/vllm-openai-rocm已存在，跳过拉取。"
    echo "现有镜像信息："
    docker images | grep "rcatrangiu/vllm-openai-rocm"
else
    echo "正在拉取Docker镜像 rcatrangiu/vllm-openai-rocm:v0.8.2 ..."
    docker pull rcatrangiu/vllm-openai-rocm:v0.8.2
    
    # 再次检查确认是否拉取成功
    if docker images | grep -q "rcatrangiu/vllm-openai-rocm.*v0.8.2"; then
        echo "Docker镜像拉取成功！"
        echo "镜像信息："
        docker images | grep "rcatrangiu/vllm-openai-rocm"
    else
        echo "错误：Docker镜像拉取失败！"
        echo "请检查："
        echo "1. Docker服务是否正常运行（systemctl status docker）"
        echo "2. 网络连接是否正常"
        echo "3. 镜像名称是否正确"
        exit 1
    fi
fi

wget https://gitee.com/huangsheng32/AMDGPU_linux/raw/master/vllm/run_vllmopenai_docker.sh
wget https://gitee.com/huangsheng32/AMDGPU_linux/raw/master/vllm/run_vllmopenai_benchmark.sh
sudo chmod +x *.sh

echo "先运行 ./run_vllmopenai_docker.sh 等待全部加载完成"
echo "再运行 ./run_vllmopenai_benchmark.sh "